package cn.edu.flink.tutorial.partition;

import cn.edu.flink.tutorial.source.UDSourceFunction;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

public class NewPartitioner {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        SingleOutputStreamOperator<Tuple2<Long, Long>> streamSource = env.addSource(new UDSourceFunction()).map(x -> Tuple2.of(x, 1L)).returns(Types.TUPLE(Types.LONG, Types.LONG));

        streamSource
                .partitionCustom(
                        new Partitioner<Long>() {
                            @Override
                            public int partition(Long key, int numPartitions) {
                                return key.intValue();
                            }
                        },
                        new KeySelector<Tuple2<Long, Long>, Long>() {
                            @Override
                            public Long getKey(Tuple2<Long, Long> value) throws Exception {
                                return value.f1;
                            }
                        }
                ).print();

        env.execute();
    }
}
